import asyncio

from langgraph.constants import START, END
from langgraph.graph import StateGraph

from xhs_auto_poster.__001__langgraph_more_nodes.agent_state import AgentState
from xhs_auto_poster.__001__langgraph_more_nodes.nodes.image_crawl_node import image_crawl_node
from xhs_auto_poster.__001__langgraph_more_nodes.nodes.intent_detector_node import intent_detector_node
from xhs_auto_poster.__001__langgraph_more_nodes.nodes.text_generate_node import text_generate_node
from xhs_auto_poster.__001__langgraph_more_nodes.nodes.image_generate_node import image_generate_node
from xhs_auto_poster.__001__langgraph_more_nodes.nodes.auto_publish_xiaohongshu_node import auto_publish_xiaohongshu_node
from xhs_auto_poster.__001__langgraph_more_nodes.nodes.llm_direct_out_node import llm_direct_out_node
from xhs_auto_poster.__001__langgraph_more_nodes.nodes.generate_markdown_node import generate_markdown_node
from xhs_auto_poster.__001__langgraph_more_nodes.nodes.check_text_image_node import check_text_image_node
from xhs_auto_poster.common.ouput_graph_utils import output_pic_graph
from xhs_auto_poster.common.logger_util import get_logger

logger = get_logger(__name__)


def build_graph():
    logger.info("开始构建LangGraph状态图")
    # 定义状态图
    graph_builder = StateGraph(AgentState)
    graph_builder.add_node(intent_detector_node.__name__, intent_detector_node)
    graph_builder.add_node(text_generate_node.__name__, text_generate_node)
    graph_builder.add_node(image_crawl_node.__name__, image_crawl_node)
    graph_builder.add_node(image_generate_node.__name__, image_generate_node)
    graph_builder.add_node(auto_publish_xiaohongshu_node.__name__, auto_publish_xiaohongshu_node)
    graph_builder.add_node(generate_markdown_node.__name__, generate_markdown_node)
    graph_builder.add_node(llm_direct_out_node.__name__, llm_direct_out_node)
    graph_builder.add_node(check_text_image_node.__name__, check_text_image_node)

    # 添加边
    graph_builder.add_edge(START, intent_detector_node.__name__)

    def intent_detector_condition(state: AgentState):
        if state['xhs_publish_intent']:
            return text_generate_node.__name__
        else:
            return llm_direct_out_node.__name__

    graph_builder.add_conditional_edges(intent_detector_node.__name__, intent_detector_condition,
                                        path_map={
                                            text_generate_node.__name__: text_generate_node.__name__,
                                            llm_direct_out_node.__name__: llm_direct_out_node.__name__,
                                        })
    graph_builder.add_edge(llm_direct_out_node.__name__, END)
    graph_builder.add_edge(text_generate_node.__name__, image_crawl_node.__name__)
    graph_builder.add_edge(image_crawl_node.__name__, check_text_image_node.__name__)

    def check_text_image_condition(state: AgentState):
        if state['is_publish_xhs']:
            return auto_publish_xiaohongshu_node.__name__
        else:
            return END

    graph_builder.add_conditional_edges(check_text_image_node.__name__, check_text_image_condition,
                                        path_map={
                                            auto_publish_xiaohongshu_node.__name__: auto_publish_xiaohongshu_node.__name__,
                                            END: END
                                        })
    graph_builder.add_edge(auto_publish_xiaohongshu_node.__name__, generate_markdown_node.__name__)
    graph_builder.add_edge(generate_markdown_node.__name__, END)

    # 编译状态图
    logger.info("编译LangGraph状态图完成")
    graph = graph_builder.compile()
    return graph


graph = build_graph()
# 输出状态图
logger.info("生成状态图可视化: graph.jpg")
output_pic_graph(graph, "graph.jpg")


async def agent_response(input: str):
    logger.info(f"接收到用户输入，开始处理: {input[:50]}..." if len(input) > 50 else f"接收到用户输入，开始处理: {input}")
    try:
        result = await graph.ainvoke({"input": input})
        logger.info("Agent处理完成")
        return result['output']
    except Exception as e:
        logger.error(f"Agent处理过程中发生错误: {str(e)}", exc_info=True)
        raise


if __name__ == '__main__':
    logger.info("启动LangGraph演示程序")
    print("请进行输入：")
    while True:
        text = input()
        if text == "q":
            logger.info("用户退出程序")
            break
        try:
            logger.debug(f"用户输入: {text}")
            response = asyncio.run(agent_response(text))
            logger.debug("生成响应完成")
        except Exception as e:
            logger.error(f"处理用户输入时出错: {str(e)}", exc_info=True)
